#!/usr/bin/python
# -*- coding: UTF-8 -*-
from __future__ import absolute_import
from __future__ import unicode_literals
import os
import cv2
import time
import traceback
import numpy as np
from skimage import morphology

def erode_op(image_input, kernel_col=3, kernel_row=3, kernel_shape='RECT'):
    kernel_shape_dict = {'RECT': cv2.MORPH_RECT,
                         'CROSS': cv2.MORPH_CROSS,
                         'ELLIPSE': cv2.MORPH_ELLIPSE
                         }
    kernel = cv2.getStructuringElement(kernel_shape_dict[kernel_shape], (kernel_col, kernel_row))
    return cv2.erode(image_input, kernel)


def dilate_op(image_input,  kernel_col=3, kernel_row=3, kernel_shape='RECT'):
    kernel_shape_dict = {'RECT': cv2.MORPH_RECT,
                         'CROSS': cv2.MORPH_CROSS,
                         'ELLIPSE': cv2.MORPH_ELLIPSE
                         }
    kernel = cv2.getStructuringElement(kernel_shape_dict[kernel_shape], (kernel_col, kernel_row))
    return cv2.dilate(image_input, kernel)


def open_op(image_input,  kernel_col=3, kernel_row=3, kernel_shape='RECT'):
    kernel_shape_dict = {'RECT': cv2.MORPH_RECT,
                         'CROSS': cv2.MORPH_CROSS,
                         'ELLIPSE': cv2.MORPH_ELLIPSE
                         }
    kernel = cv2.getStructuringElement(kernel_shape_dict[kernel_shape], (kernel_col, kernel_row))  # 提取字母
    return cv2.morphologyEx(image_input, cv2.MORPH_OPEN, kernel)


def close_op(image_input,  kernel_col=3, kernel_row=3, kernel_shape='RECT'):
    kernel_shape_dict = {'RECT': cv2.MORPH_RECT,
                         'CROSS': cv2.MORPH_CROSS,
                         'ELLIPSE': cv2.MORPH_ELLIPSE
                         }
    kernel = cv2.getStructuringElement(kernel_shape_dict[kernel_shape], (kernel_col, kernel_row))  # 提取字母
    return cv2.morphologyEx(image_input, cv2.MORPH_CLOSE, kernel)


def skeleton_extract(image_input, bw_th=128, method=1):
    """
    说明：
    这种算法能将一个连通区域细化成一个像素的宽度，用于特征提取和目标拓扑表示。骨架提取与分水岭算法也属于形态学处理范畴
    morphology子模块提供了两个函数用于骨架提取，分别是Skeletonize()函数和medial_axis()函数
    ( medial_axis就是中轴的意思，利用中轴变换方法计算前景（1值）目标对象的宽度 )
    """
    if image_input.ndim == 3:
        image_2d = cv2.cvtColor(image_input, cv2.COLOR_BGR2GRAY)
    else:
        image_2d = np.array(image_input)
    _, image_binary = cv2.threshold(image_2d, bw_th, 255, cv2.THRESH_BINARY)
    image_binary[image_binary == 255] = 1

    if method == 1:
        image_skeleton = morphology.skeletonize(image_binary)
        image_skeleton = image_skeleton.astype(np.uint8) * 255
        return image_skeleton
    elif method == 2:
        skel, distance = morphology.medial_axis(image_binary, return_distance=True)
        image_skeleton = distance * skel
        image_skeleton = image_skeleton.astype(np.uint8) * 255
        return image_skeleton



if __name__ == "__main__":
    path = r'C:\Users\admin\Desktop\dt\test.bmp'  # place path to your image here
    image_cv = cv2.imread(path, -1)
    # cv2.imshow("original", image_cv)
    # cv2.waitKey(0)
    gau = skeleton_extract(image_cv, 120, 2)
    cv2.imshow("Operation", gau)
    cv2.waitKey(0)